import os
import torch
from torch import nn
from dataload import get_dataload
from util import try_gpu

class Trainer:
    def __init__(self, model, train_evaluator, test_evaluator, lr, num_epoch, batch_size):
        self.model = model
        self.train_evaluator = train_evaluator
        self.test_evaluator = test_evaluator
        self.num_epoch = num_epoch
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=lr, weight_decay=5e-4, momentum=0.9)
        self.loss = nn.CrossEntropyLoss()
        self.best_acc = 0
        self.dataloader = get_dataload("./BigData/dataset", my_batch_size=batch_size, is_train=True)

    def train(self):
        output_dir = "./Output/"
        if not os.path.exists(output_dir):
            os.mkdir(output_dir)
        file = open(output_dir + "train_log.txt", 'w')
        for epoch in range(self.num_epoch):
            self.model.train()
            total_l = 0
            avg_l = 0
            for i, (input, label) in enumerate(self.dataloader):
                input = input.to(device=try_gpu())
                label = label.to(device=try_gpu())
                output = self.model(input)
                l = self.loss(output, label)
                total_l += l.item()
                self.optimizer.zero_grad()
                l.backward()
                self.optimizer.step()
                avg_l = total_l / (i + 1)
                if (i+1) % 10 == 0:
                    print(i+1, " finish! avg_l: ", avg_l)
            print("epoch: ", epoch, "   loss: ", avg_l)
            file.write("epoch: " + str(epoch) + "   loss: " + str(avg_l) + "\n")
            if (epoch + 1) % 2 == 0:
                train_accuracy = self.train_evaluator.eval(self.model)
                test_accuracy = self.test_evaluator.eval(self.model)
                print("\ntrain_acc: ", train_accuracy, "   test_acc: ", test_accuracy, "\n")
                file.write("\ntrain_acc: " + str(train_accuracy) + "   test_acc: " + str(test_accuracy) + "\n")
                if test_accuracy > self.best_acc:
                    torch.save(self.model, output_dir + "best.pt")
                    self.best_acc = test_accuracy
                    print("-"*10, "save model!", "-"*10)
                    
                for params in self.optimizer.param_groups:
                    params['lr'] = params['lr'] - 0.1 * (params['lr'] - 0.0001)
                    print("lr: ", params['lr'])